"""Class for handling LD matrices"""

import numpy as np


EPS = np.finfo(float).eps
EPS2 = EPS ** 2

def ltdl(mat_x):
    """
    L'DL decomposition
    
    Parameters
    ----------
    max_x : numpy.array
           Matrix to be L'DL factorized

    Returns
    -------
    mat_l : numpy.array
           Unit lower triangular matrix
    vec_d : numpy.array
           Vector containing diagonal entries of D matrix
    """
    mat_x = mat_x[::-1, ::-1]
    mat_x = np.linalg.cholesky(mat_x)
    mat_x = mat_x[::-1, ::-1].T
    mat_x_diag = np.diag(mat_x)
    vec_d = mat_x_diag ** 2

    for i in xrange(1, len(mat_x_diag)):
        s = 1/mat_x_diag[i] if (mat_x_diag[i] > EPS) else 0
        mat_x[i] *= s

    mat_l = mat_x - np.diag(np.diag(mat_x)) + np.eye(mat_x.shape[0])

    return mat_l, vec_d


def dydrs(vec_data, row_of_mat_l, weight, d_j, R, jl, jh):
    """
    Diadic reduction

    Stupid rewrite of dydrs from mixtools.

    Parameters
    ----------
    d_j : float
         j-th element of D matrix
    vec_data : numpy.array
         vector with data
    row_of_mat_l : int
        which row of L is changed
    weight : float
        weight of update
    R : int
        order of data
    jl : int
        lower coordinate j of L matrix
    jh : int
        higher coordinate j of L matrix

    Returns
    -------
    vec_data_new : numpy.array
        Vector on new data
    row_of_mat_l_new : numpy.array
        Vector of new L row
    weights_new : float
        New weight
    d_j_new : float
        New D_j
    """

    row_of_mat_l[R] = d_j

    jh += 1

    if np.abs(weight) < EPS2:
        weight = 0

    data_item = vec_data[R]
    kD = d_j
    kr = data_item * weight
    d_j_new = kD + data_item * kr

    if d_j_new > EPS2:
        kD = kD / d_j_new
        kr = kr / d_j_new
    else:
        kD = 1
        kr = 0
        d_j_new = 0

    weight_new = weight * kD
    vec_data_new = vec_data.copy()
    row_of_mat_l_new = row_of_mat_l.copy()
    vec_data_new[jl:jh] = vec_data[jl:jh] - data_item * row_of_mat_l[jl:jh]
    row_of_mat_l_new[jl:jh] = row_of_mat_l[jl:jh] + kr * vec_data_new[jl:jh]
    row_of_mat_l_new[R] = 1


    return vec_data_new, row_of_mat_l_new, weight_new, d_j_new


class Ldmat(object):
    """
    LD representation of information matrix

    Attributes
    ----------
    _l : numpy.array
        Lower triangular matrix
    _d : numpy.array
        Vector of the diagonal matrix
    _dim: int
        Dimension of the matrix


    Methods
    -------
    __init__(mat_v, from_ld=False)
        Constructor
    update(data)
        Sherman-Morrison rank-1 update
    perm(perm_vec)
        Permutation of rows
    ldform(A, D0)
        Form LD matrix from arbitrary matrix and vector D0


    Properties
    ----------
    v : getter for full matrix
    l, d : getters for L and D
    ld : getter for concatenated form LD
    shape : getter for matrix shape
    """


    def __init__(self, mat_v, from_ld=False):
        """
        Constructor

        Parameters
        ----------
        mat_v : numpy.array
            Information matrix
        from_ld : bool
            True if mat_v is instance of Ldmat 
        """
        #We silently suppose that mat_v is square
        self._dim = mat_v.shape[0]
        if not from_ld:
            self._l, self._d = ltdl(mat_v)
        elif from_ld == True:
            self._d = np.diag(mat_v)
            self._l = mat_v - np.diag(self._d) + np.eye(self._dim)
        
    @property
    def l(self):
        """Return L' matrix"""
        return self._l

    @property
    def d(self):
        """Return D vector"""
        return self._d

    @property
    def v(self):
        """Return V matrix"""
        mat_v = np.dot(self._l.T, np.diag(self._d))
        mat_v = np.dot(mat_v, self._l)
        return mat_v

    @property
    def ld(self):
        """Return matrix in `concatenate' form"""
        return self._l - np.eye(self._dim) + np.diag(self._d)

    @property
    def shape(self):
        """Shape of LD matrix"""
        return self._l.shape

    def update(self, data):
        """
        Rank-one update of matrix in LD form
        
        Parameters
        ----------
        data : numpy.array
            Vector of data
        """
        if len(data) != self._dim:
            raise ValueError\
            ('Dimension of data doesn\'t match size of inf. matrix!')

        weight = 1.

        for j in xrange(self._dim - 1, -1, -1):
            if data[j] != 0:
                l_tmp = self._l[j]
                [data, self._l[j], weight, d_j_new] = \
                    dydrs(data, self.l[j], weight, self._d[j], j, 0, j-1)
                self._d[j] = d_j_new


    def ldform(self, A, D0):
        """
        Form LD matrix from general matrix A and vector D
        Copied from Vasek Smidl's code.
        """
        m, n = A.shape
        mn = np.min((m, n))
    
        arg_tmp = np.dot(np.diag(np.sqrt(D0)), A)
        L = np.concatenate((np.zeros((n, n)), arg_tmp))
        D = np.zeros(n + m)
    
        w = np.zeros(n + m)
    
        cc = 0
        i = n
        while (i > n - mn - cc) and (i > 0): 
            i-=1
            sum = 0.
            last_v = m + i - n + cc + 1
    
            v = np.zeros(last_v + 1)
            for ii in xrange(n - cc - 1, m + i + 1): 
                sum += L[ii, i] * L[ii, i]
                v[ii - n + cc + 1] = L[ii, i]
    
            if  L[m + i, i] == 0:
                beta = np.sqrt(sum)
            else:
                beta = L[m + i, i] + np.sign(L[m + i, i]) * np.sqrt(sum);
    
            if (np.abs(beta) < np.finfo(float).eps ):
                print "Ldmat: Numerical difficulties!!!"
                cc += 1
                L[n - cc] = L[m + i]
                L[m + i] *= 0
                D[m + i] = 0
                L[m + i, i] = 1
                L[n - cc:m + i, i:i + 1] *= 0
                continue
    
            sum -= v[last_v] ** 2
            sum /= beta ** 2
            sum += 1
    
            v /= beta;
            v[last_v] = 1
    
            pom = -2.0 / sum
    
            for j in xrange(i, -1, -1):
                w_elem = 0
                for ii in xrange(n - cc, m + i + 2):
                    w_elem += np.dot(v[ii - n + cc], L[ii - 1, j])
                w[j] = w_elem * pom
    
            for ii in xrange(n - cc - 1, m + i + 1):
                for jj in xrange(0, i):
                    L[ii, jj] += np.dot(v[ii - n + cc + 1], w[jj])
    
            for ii in xrange(n - cc - 1, m + i):
                L[ii, i] = 0
    
            L[m + i, i] += w[i]
            D[m + i] = np.dot(L[m + i, i], L[m + i, i])
    
            for ii in xrange(0, i + 1):
                L[m + i, ii] /= L[m + i, i]
    
        if i >= 0:
            for ii in xrange(0, i):
                jj = D.shape[0] - 1 - n + ii
                D[jj] = 0
                L[jj] *= 0
                L[jj, jj] = 1
    
        self._l = L[m:];
        self._d = D[m:];
   

    def perm(self, perm_vec):
        """
        Permutation of rows in L'DL factorized matrix

        Parameters
        ----------
        perm_vec : numpy.array
            Permutation vectors with new order indices of rows
            e.g. [3, 2, 1] reverses the rows of L'DL.
        """
        self.ldform(self._l[:, perm_vec], self._d)
    


# =============================================================
if __name__ == "__main__":
    print "Ldmat test"

    V = np.array([[5., 4., 3., 2., 1.],
                  [0., 6., 1., 1., 2.],
                  [0., 0., 7., 2., 3.], 
                  [0., 0., 0., 6., 2.], 
                  [0., 0., 0., 0., 3.]])
    V = V + V.T
    print V

    LD = Ldmat(V)

    dt = np.array([1., 1., 1., 1., 1.])
    
#    vd = np.array([1., 1., 1., 1.])
#    rml = np.array([1., 0, 0, 0])
#    dj = 4.
#    print dydrs(vd, rml, 1, dj, 1, 1, 2)
    
    LD.update(dt)
    print V - LD.v
